package com.service.data.flume.tools.rabbitmq.sinks

import com.rabbitmq.client.{Channel, Connection, ConnectionFactory}
import com.service.data.commons.utils.Preconditions
import org.apache.flume.Sink.Status
import org.apache.flume.{Context, Event, Sink}
import org.apache.flume.conf.Configurable
import org.apache.flume.sink.AbstractSink

/**
  * @author 伍鲜
  *
  *         Flume Sink to RabbitMQ
  */
class RabbitMQSink extends AbstractSink with Configurable {

  private var factory: ConnectionFactory = _
  private var connection: Connection = _
  private var rabbitChannel: Channel = _

  private var hostname: String = _
  private var username: String = _
  private var password: String = _
  private var port: Int = _

  private var exchange: String = _
  private var routingKey: String = _
  private var mandatory: Boolean = _
  private var immediate: Boolean = _

  /**
    * 配置，环境初始化
    *
    * @param context
    */
  override def configure(context: Context): Unit = {
    // 必选项
    hostname = Preconditions.checkNotNull(context.getString("hostname"), "Rabbit MQ hostname must be specified. For example: agent.sinks.sink.hostname=127.0.0.1")
    exchange = Preconditions.checkNotNull(context.getString("exchange"), "Rabbit MQ exchange must be specified. For example: agent.sinks.sink.exchange=wux.exchange")

    // 可选项
    username = context.getString("username", null)
    password = context.getString("password", null)
    port = context.getInteger("port", -1)
    routingKey = context.getString("routingKey", "")
    mandatory = context.getBoolean("mandatory", false)
    immediate = context.getBoolean("immediate", false)

    factory = new ConnectionFactory
    factory.setHost(hostname)

    if (username != null) factory.setUsername(username)
    if (password != null) factory.setPassword(password)
    if (port != -1) factory.setPort(port)

    connection = factory.newConnection
    rabbitChannel = connection.createChannel
  }

  override def process(): Sink.Status = {
    val flumeChannel = getChannel
    val tx = flumeChannel.getTransaction
    try {
      // 事务开始
      tx.begin()
      val event: Event = flumeChannel.take
      if (!(event == null || event.getBody == null || (event.getBody.length == 0))) {
        rabbitChannel.basicPublish(exchange, routingKey, mandatory, immediate, null, event.getBody)
      }
      // 事务提交
      tx.commit()
      Status.READY
    } catch {
      case _: Exception =>
        try {
          // 事务回滚
          tx.rollback()
        } catch {
          case _: Exception =>
        }
        Status.BACKOFF
    } finally {
      // 事务关闭
      tx.close()
    }
  }

  override def stop() = {
    try {
      rabbitChannel.close()
    } catch {
      case _: Exception =>
    }
    super.stop()
  }
}